package com.tjetc.controller.messages;


import com.tjetc.config.OpenAiConfig;
import com.tjetc.model.AssistantMessage;
import com.tjetc.model.ChatResponse;
import com.tjetc.model.Message;
import com.tjetc.model.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class OpenAiChatModel {
    private final WebClient webClient;
    private final OpenAiConfig openAiConfig;

    @Autowired
    public OpenAiChatModel(WebClient webClient, OpenAiConfig openAiConfig) {
        this.webClient = webClient;
        this.openAiConfig = openAiConfig;
    }

    public ChatResponse call(Prompt prompt) {
        // 使用配置的baseUrl构建完整的API URL
        String apiUrl = openAiConfig.getBaseUrl() + "/chat/completions";

        // 构建请求体
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", openAiConfig.getChat().getOptions().getModel());
        requestBody.put("messages", prompt.getMessages().stream()
                .map(message -> {
                    Map<String, String> messageMap = new HashMap<>();
                    messageMap.put("role", message.getRole());
                    messageMap.put("content", message.getContent());
                    return messageMap;
                })
                .collect(Collectors.toList()));

        // 使用WebClient发送请求
        Map response = webClient.post()
                .uri(apiUrl)
                .contentType(MediaType.APPLICATION_JSON)
                .header("Authorization", "Bearer " + openAiConfig.getApiKey())
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(Map.class)
                .block(); // 虽然使用block()会阻塞，但保持API兼容性

        // 解析响应并构建ChatResponse对象
        ChatResponse chatResponse = new ChatResponse();
        ChatResponse.Result result = new ChatResponse.Result();

        if (response != null && response.containsKey("choices")) {
            List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");

            if (!choices.isEmpty()) {
                Map<String, Object> choice = choices.get(0);
                Map<String, String> messageMap = (Map<String, String>) choice.get("message");

                Message assistantMessage = new AssistantMessage(messageMap.get("content"));
                result.setOutput(assistantMessage);

                // 设置其他元数据
                ChatResponse.Metadata metadata = new ChatResponse.Metadata();
                metadata.setFinishReason((String) choice.get("finish_reason"));
                result.setMetadata(metadata);
                result.setEmpty(false);
            }
        }

        chatResponse.setResult(result);
        return chatResponse;
    }
}